#!/usr/bin/env python3
"""
Advanced Access Control List Controller
Implementing ACL with P4Runtime
Features:
1. Drop packets with dstIP=10.0.1.4
2. Drop UDP packets with dstPort=80
"""

import argparse
import grpc
import os
import sys
from time import sleep

# Import P4Runtime lib from parent utils dir
sys.path.append(os.path.join(os.path.dirname(os.path.abspath(__file__)), "../utils/"))
import p4runtime_lib.bmv2
from p4runtime_lib.error_utils import printGrpcError
from p4runtime_lib.switch import ShutdownAllSwitchConnections
import p4runtime_lib.helper


def writeIpv4ForwardRules(p4info_helper, sw, match_fields, action_params):
    """Write IPv4 forwarding rules to switch"""
    table_entry = p4info_helper.buildTableEntry(
        table_name="MyIngress.ipv4_lpm",
        match_fields={"hdr.ipv4.dstAddr": match_fields},
        action_name="MyIngress.ipv4_forward",
        action_params=action_params,
    )
    sw.WriteTableEntry(table_entry)
    print("Added IPv4 forward rule: %s" % match_fields[0])


def writeAclIpRules(p4info_helper, sw, match_fields, priority):
    """Write IP address ACL rules to switch"""
    table_entry = p4info_helper.buildTableEntry(
        table_name="MyIngress.acl_ip_ternary",
        match_fields={"hdr.ipv4.dstAddr": match_fields},
        action_name="MyIngress.drop",
        priority=priority,
    )
    sw.WriteTableEntry(table_entry)
    print("Added IP ACL rule to drop: %s" % match_fields[0])


def writeAclUdpRules(p4info_helper, sw, match_fields, priority):
    """Write UDP port ACL rules to switch"""
    table_entry = p4info_helper.buildTableEntry(
        table_name="MyIngress.acl_udp_ternary",
        match_fields={"hdr.udp.dstPort": match_fields},
        action_name="MyIngress.drop",
        priority=priority,
    )
    sw.WriteTableEntry(table_entry)
    print("Added UDP ACL rule to drop port: %s" % match_fields[0])


def setupSwitchRules(p4info_helper, switch):
    """Setup all necessary rules on the switch"""
    print("Setting up switch rules...")

    # IPv4 forwarding rules for all hosts
    writeIpv4ForwardRules(
        p4info_helper,
        switch,
        ["10.0.1.1", 32],
        {"dstAddr": "08:00:00:00:01:01", "port": 1},
    )
    writeIpv4ForwardRules(
        p4info_helper,
        switch,
        ["10.0.1.2", 32],
        {"dstAddr": "08:00:00:00:01:02", "port": 2},
    )
    writeIpv4ForwardRules(
        p4info_helper,
        switch,
        ["10.0.1.3", 32],
        {"dstAddr": "08:00:00:00:01:03", "port": 3},
    )
    writeIpv4ForwardRules(
        p4info_helper,
        switch,
        ["10.0.1.4", 32],
        {"dstAddr": "08:00:00:00:01:04", "port": 4},
    )

    # ACL Rule 1: Drop packets with destination IP 10.0.1.4
    writeAclIpRules(p4info_helper, switch, ["10.0.1.4", 0xFFFFFFFF], 10)

    # ACL Rule 2: Drop UDP packets with destination port 80
    writeAclUdpRules(p4info_helper, switch, [80, 0xFFFF], 10)

    print("All rules configured successfully!")


def main(p4info_file_path, bmv2_file_path):
    """Main function to setup switch and configure rules"""
    # Initialize P4Info helper
    p4info_helper = p4runtime_lib.helper.P4InfoHelper(p4info_file_path)

    try:
        # Create switch connection
        print("Connecting to switch s1...")
        s1 = p4runtime_lib.bmv2.Bmv2SwitchConnection(
            name="s1",
            address="127.0.0.1:50051",
            device_id=0,
            proto_dump_file="logs/s1-p4runtime-requests.txt",
        )

        # Establish master arbitration
        print("Establishing master arbitration...")
        s1.MasterArbitrationUpdate()

        # Set forwarding pipeline config
        print("Setting forwarding pipeline config...")
        s1.SetForwardingPipelineConfig(
            p4info=p4info_helper.p4info, bmv2_json_file_path=bmv2_file_path
        )

        # Setup all rules
        setupSwitchRules(p4info_helper, s1)

        print("ACL configuration completed!")
        print("Rules summary:")
        print("1. Drop all packets to IP 10.0.1.4")
        print("2. Drop all UDP packets to port 80")
        print("3. Allow other traffic with IPv4 forwarding")

    except KeyboardInterrupt:
        print("\nShutting down.")
    except grpc.RpcError as e:
        printGrpcError(e)
    finally:
        ShutdownAllSwitchConnections()


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="P4Runtime ACL Controller")
    parser.add_argument(
        "--p4info",
        help="p4info proto in text format from p4c",
        type=str,
        action="store",
        required=False,
        default="./build/acl.p4.p4info.txt",
    )
    parser.add_argument(
        "--bmv2-json",
        help="BMv2 JSON file from p4c",
        type=str,
        action="store",
        required=False,
        default="./build/acl.json",
    )
    args = parser.parse_args()

    # Validate file existence
    if not os.path.exists(args.p4info):
        parser.print_help()
        print("\nError: p4info file not found: %s" % args.p4info)
        print("Please run 'make' to compile the P4 program first.")
        parser.exit(1)
    if not os.path.exists(args.bmv2_json):
        parser.print_help()
        print("\nError: BMv2 JSON file not found: %s" % args.bmv2_json)
        print("Please run 'make' to compile the P4 program first.")
        parser.exit(1)

    # Run main function
    main(args.p4info, args.bmv2_json)
